from PIL import Image
import os
import time
import cv2
import numpy as np


class PSTools:

    @staticmethod
    def mask_expander(mask_path, out_path=None, pixels_to_expand=2):
        """蒙版外扩

        Args:
            mask_path (_type_): _description_
            out_path (_type_, optional): _description_. Defaults to None.
            pixels_to_expand (int, optional): _description_. Defaults to 2.

        Returns:
            _type_: _description_
        """
        # 加载蒙版图
        mask = cv2.imread(mask_path, 0)
        mask = cv2.bitwise_not(mask)
        # 创建结构元素
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        # 扩展黑色区域
        dilated_mask = cv2.dilate(mask, kernel, iterations=pixels_to_expand)
        dilated_mask = cv2.bitwise_not(dilated_mask)
        if out_path:
            cv2.imwrite(out_path, dilated_mask)
        return dilated_mask

    @staticmethod
    def find_max_mask(image_path, out_path=None, bitwise_not=True):
        """找蒙版中得最大区域

        Args:
            image_path (_type_): _description_
            out_path (_type_, optional): _description_. Defaults to None.
            bitwise_not (bool, optional): _description_. Defaults to True.

        Returns:
            _type_: _description_
        """
        # 加载图像
        image = cv2.imread(image_path, 0)  # 加载为灰度图像
        if bitwise_not:
            image = cv2.bitwise_not(image)
        # 分割区域
        ret, binary_image = cv2.threshold(
            image, 127, 255, cv2.THRESH_BINARY)  # 二值化图像
        # 查找轮廓
        contours, hierarchy = cv2.findContours(
            binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # 计算面积
        areas = [cv2.contourArea(cnt) for cnt in contours]
        # 选择最大区域
        max_area_index = np.argmax(areas)
        max_area_contour = contours[max_area_index]
        # 创建一个与原图像大小相同的掩膜图像
        mask = np.zeros_like(image)
        # 绘制最大区域的轮廓
        cv2.drawContours(mask, [max_area_contour], 0, 255, -1)
        # 将小区域置为白色
        result = cv2.bitwise_and(image, mask)
        if bitwise_not:
            result = cv2.bitwise_not(result)
        if out_path:
            cv2.imwrite(out_path, result)
        return result

    @staticmethod
    def mask_bitwise_not(mask_path, out_path=None):
        mask = cv2.imread(mask_path, 0)
        mask = cv2.bitwise_not(mask)
        if out_path:
            cv2.imwrite(out_path, mask)
        return mask

    @staticmethod
    def mask_remove_isolated_pixels(mask_path, out_path=None, open_kernel_size=3, close_kernel_size=5):
        """移除蒙版中得孤立点

        Args:
            mask_path (_type_): _description_
            out_path (_type_, optional): _description_. Defaults to None.
            open_kernel_size (int, optional): _description_. Defaults to 3.
            close_kernel_size (int, optional): _description_. Defaults to 5.

        Returns:
            _type_: _description_
        """
        image = cv2.imread(mask_path, 0)
        # 创建一个结构元素
        kernel = cv2.getStructuringElement(
            cv2.MORPH_RECT, (open_kernel_size, open_kernel_size))
        # 开运算，移除孤立点
        image = cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)
        # 创建一个结构元素
        kernel = cv2.getStructuringElement(
            cv2.MORPH_RECT, (close_kernel_size, close_kernel_size))
        # 闭运算，填补小的空洞
        image = cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel)
        if out_path:
            cv2.imwrite(out_path, image)
        return image

    @staticmethod
    def compare_mask_and_origin(image_path, mask_path, save_dir, background_color=(0, 0, 0, 0)):
        """蒙版与原图比较，得到原图中得像素值

        Args:
            image_path (_type_): _description_
            mask_path (_type_): _description_
            save_dir (_type_): _description_
            background_color (tuple, optional): _description_. Defaults to (0, 0, 0, 0).

        Returns:
            _type_: _description_
        """
        original_image = Image.open(image_path)
        mask_image = Image.open(mask_path)
        mask_image = mask_image.convert("L")
        white_background = Image.new(
            "RGBA", original_image.size, (255, 255, 255))

        # 创建一个与原始图像大小相同的白色背景图像
        blended_image = Image.new("RGBA", original_image.size, (255, 255, 255))
        original_pixels = original_image.load()
        blended_pixels = blended_image.load()
        mask_pixels = mask_image.load()

        width, height = original_image.size
        for x in range(width):
            for y in range(height):
                if mask_pixels[x, y] == 0:  # 将判断条件修改为0，即黑色像素值
                    blended_pixels[x, y] = original_pixels[x, y]  # 使用原始图像的像素值
                else:
                    blended_pixels[x, y] = background_color  # 其他位置设置为白色背景
        filename = os.path.join(save_dir, "body-seg-"+str(time.time())+".png")
        blended_image.save(filename)
        return filename

    @staticmethod
    def merge_mask(mask_paths=[], out_path=""):
        """合并多个蒙版图,注意蒙版之间加法运算

        Args:
            mask_paths (list, optional): _description_. Defaults to [].
            out_path (str, optional): _description_. Defaults to "".

        Returns:
            _type_: _description_
        """
        init_image = cv2.imread(mask_paths[0], 0)
        final_image = init_image
        for path_ in mask_paths:
            _image = cv2.imread(path_, 0)
            final_image += cv2.resize(_image,
                                      (init_image.shape[1], init_image.shape[0]))
        mask = np.clip(final_image, 0, 255)
        if out_path:
            cv2.imwrite(out_path, mask)
        return mask


class PSCanvasLayer:
    def __init__(self, background_path=None, canvas_width=512, canvas_height=680):
        if background_path:
            self.background = Image.open(background_path).convert('RGBA')
        else:
            self.background = Image.new(
                'RGBA', (canvas_width, canvas_height), (255, 255, 255, 255))
        self.layers = []

    def add_layer(self, layer_path, position):
        layer = Image.open(layer_path).convert('RGBA')
        self.layers.append((layer, position))

    def merge_layers(self, save_dir):
        merged_image = self.background.copy()
        for layer, position in self.layers:
            merged_image = self._merge_layer(merged_image, layer, position)
        filename = os.path.join(save_dir, "ps-"+str(time.time())+".png")
        merged_image.save(filename)
        return filename

    def _merge_layer(self, background, layer, position):
        merged = Image.new('RGBA', background.size)
        merged.paste(background, (0, 0))
        merged.paste(layer, position, mask=layer)
        return merged
